import tools
import time

class Nsga_ii:
    """
    基础版本的NSGA-II算法
    """
    def __init__(self, change_ratio, max_train_time, city_number, gene_number):
        self.change_ratio=change_ratio #交叉变异率
        self.max_train_time=max_train_time #最大训练时间（分钟）

        city_ids, customer, time_matrix, _ = tools.generate_city_graph(n=city_number, show_img=False)
        self.city_ids=city_ids #每个城市的编号
        self.city_customers = customer #每个城市有多少人
        self.time_matrix=time_matrix #各个城市之间的运输时间(矩阵)

        self.gene_number = gene_number
        self.current_genes = None
        #current_genes是一个列表，其中每个元素是一个基因。

        self.fitness = None #记录每个基因在两个目标上的表现

        self.cm_number = 40

    def get_fitness(self): #计算当前种群在各个目标上的表现
        self.fitness = tools.calculate_fitness(population=self.current_genes, customer=self.city_customers,
                                               time_matrix=self.time_matrix)

    def get_new_cm_number(self):
        """
        基于变异率重新计算每次选择多少个基因交叉变异。
        """
        self.cm_number = int(self.change_ratio*len(self.current_genes))

        if self.cm_number%2==1:
            self.cm_number+=1


    def find_result(self):
        """
        用来求解优秀路径的函数
        """
        start_time = time.time()

        self.current_genes = tools.initialize_population(cities=self.city_ids, N=self.gene_number) # 嵌套列表

        best_result_recorder=[] #用来记录最上层基因在两个目标上的平均表现


        while True:

            self.get_fitness()  # 计算这批基因在各个指标上表现，fitness是嵌套列表
            fronts = tools.non_dominated_sort(self.current_genes,self.fitness) #嵌套列表

            best_front=fronts[0] #最好基因的 下标

            cost=0
            awt = 0
            for index in best_front:
                awt += self.fitness[index][0]
                cost += self.fitness[index][1]

            best_result_recorder.append([awt/len(best_front),cost/len(best_front)])


            # 计算当前循环运行了多少秒
            elapsed_time = time.time() - start_time
            # 将秒转换为分钟
            elapsed_minutes = elapsed_time / 60

            if elapsed_minutes > self.max_train_time:
                print(f'训练时间超过{self.max_train_time}分钟，结束运行。')
                break

            good_genes = tools.select_best_genes(self.current_genes,self.fitness,fronts,n=self.cm_number) #用来交叉排序的优秀基因
            child = tools.crossover_and_mutation(good_genes) #交叉排序后得到的子代基因。

            self.current_genes.extend(child)

            self.get_fitness()
            fronts = tools.non_dominated_sort(self.current_genes, self.fitness)
            self.current_genes = tools.select_best_genes(self.current_genes,self.fitness,fronts,n=self.gene_number) #得到下一代基因

        all_awt = []
        all_cost = []

        for item in best_result_recorder:
            all_awt.append(item[0])
            all_cost.append(item[1])

        tools.plot_subgraphs(all_awt,all_cost)


class Nsga_ii_version2:
    """
    改进版本的NSGA-II算法
    """
    def __init__(self, change_ratio, max_train_time,city_number, gene_number):
        self.change_ratio=change_ratio #刚开始时候的交叉变异率
        self.max_train_time=max_train_time #最大训练时间（分钟）

        city_ids, customer, time_matrix, _ = tools.generate_city_graph(n=city_number, show_img=False)
        self.city_ids=city_ids #每个城市的编号
        self.city_customers = customer #每个城市有多少人
        self.time_matrix=time_matrix

        self.gene_number = gene_number
        self.current_genes = None
        #current_genes是一个列表，其中每个元素是一个基因。

        self.fitness = None #记录每个基因在两个目标上的表现

        self.cm_number = 140

    def get_fitness(self): #计算当前种群在各个目标上的表现
        self.fitness = tools.calculate_fitness(population=self.current_genes, customer=self.city_customers,
                                               time_matrix=self.time_matrix)

    def get_new_cm_number(self):
        """
        基于变异率重新计算每次选择多少个基因交叉变异。
        """
        self.cm_number = int(self.change_ratio*len(self.current_genes))

        if self.cm_number%2==1:
            self.cm_number+=1


    def find_result(self):
        """
        用来求解优秀路径的函数
        """
        start_time = time.time()

        self.current_genes = tools.initialize_population_version2(cities=self.city_ids, customers=self.city_customers,
                                                                  time_matrix=self.time_matrix, N=self.gene_number)


        best_result_recorder=[] #用来记录最上层基因在两个目标上的平均表现

        while True:

            self.get_fitness()  # 计算这批基因在各个指标上表现，fitness是嵌套列表
            fronts = tools.non_dominated_sort(self.current_genes,self.fitness) #嵌套列表

            best_front=fronts[0] #最好基因的 下标

            cost = 0
            awt = 0
            for index in best_front:
                awt += self.fitness[index][0]
                cost += self.fitness[index][1]

            best_result_recorder.append([awt / len(best_front), cost / len(best_front)])


            # 计算当前循环运行了多少秒
            elapsed_time = time.time() - start_time
            # 将秒转换为分钟
            elapsed_minutes = elapsed_time / 60

            if elapsed_minutes > self.max_train_time:
                print(f'训练时间超过{self.max_train_time}分钟，结束运行。')
                break


            good_genes = tools.select_best_genes(self.current_genes,self.fitness,fronts,n=self.cm_number) #用来交叉排序的优秀基因
            child = tools.crossover_and_mutation(good_genes) #交叉排序后得到的子代基因。

            self.current_genes.extend(child)

            self.get_fitness()
            fronts = tools.non_dominated_sort(self.current_genes, self.fitness)
            self.current_genes = tools.select_best_genes(self.current_genes,self.fitness,fronts,n=self.gene_number) #得到下一代基因

            temp = self.change_ratio*0.97
            if temp>0.05:
                self.change_ratio=temp
                self.get_new_cm_number()

        all_awt = []
        all_cost = []

        for item in best_result_recorder:
            all_awt.append(item[0])
            all_cost.append(item[1])

        tools.plot_subgraphs(all_awt,all_cost)


if __name__ == '__main__':
    max_train_time = 1
    city_numbers = 20
    gene_number = 200

    city_ids, customer, distance_matrix, G = tools.generate_city_graph(n=city_numbers,show_img=False)


    nsga = Nsga_ii(change_ratio=0.2, max_train_time=max_train_time, city_number=city_numbers, gene_number=gene_number)

    nsga.city_ids = city_ids
    nsga.city_customers = customer
    nsga.time_matrix = distance_matrix

    nsga_v2 = Nsga_ii_version2(change_ratio=0.7, max_train_time=max_train_time, city_number=city_numbers,
                               gene_number=gene_number)

    nsga_v2.city_ids = city_ids
    nsga_v2.city_customers = customer
    nsga_v2.time_matrix = distance_matrix

    for _ in range(3):
        print('===================================nsga_V2======================================')
        nsga_v2.find_result()

        print('===================================nsga======================================')
        nsga.find_result()







